import codecs
import csv
import os
import numpy as np
import cv2

import data.Datasets.CASIA.split_train_test as casia_util
import data.Datasets.IITD.transform_IITD as iitd_util


class DataLoader(object):
    def __init__(self, image_path):
        self._img_path = image_path

    def map(self, mapper, *args, **kwargs):
        pass

    def get(self, filename):
        pass


class CASIALoader(DataLoader):
    def map(self, mapper, *args, **kwargs):
        dirs = os.listdir(self._img_path)
        path_save, rio_extractor, writer = args[0], args[1], args[2]
        M_Time = []
        for dir in dirs:
            if not os.path.isdir(os.path.join(self._img_path, dir)):
                continue
            files = os.listdir(os.path.join(self._img_path, dir))
            for ifile in files:
                if ifile == 'Thumbs.db':
                    continue
                img = cv2.imread(os.path.join(self._img_path, dir, ifile), cv2.IMREAD_COLOR)
                timecost = mapper(img, path_save, rio_extractor, writer, ifile)
                M_Time.append(timecost)
        return np.array(M_Time)

    def get(self, filename):
        dir = casia_util.get_pid(filename)
        img = cv2.imread(os.path.join(self._img_path, str(dir), filename), cv2.IMREAD_COLOR)
        return img


class CASIANormalLoader(DataLoader):
    def map(self, mapper, *args, **kwargs):
        files = os.listdir(self._img_path)
        M_Time = []
        for ifile in files:
            if ifile == 'Thumbs.db':
                continue
            if ifile.find('_r_') == -1:
                continue
            img = cv2.imread(os.path.join(self._img_path, ifile), cv2.IMREAD_GRAYSCALE)
            timecost = mapper(img, ifile, *args, **kwargs)
            M_Time.append(timecost)
        return np.array(M_Time)

    def get(self, filename):
        dir = casia_util.get_pid(filename)
        img = cv2.imread(os.path.join(self._img_path, str(dir), filename), cv2.IMREAD_COLOR)
        return img

class IITDLoader(DataLoader):
    def __init__(self, img_path):
        super(img_path)

    def __init__(self, img_path, label_csv):
        super(IITDLoader, self).__init__(img_path)
        # super(DataLoader, img_path)
        self.csv_path = label_csv

    def map(self, mapper, *args, **kwargs):
        files = os.listdir(self._img_path)
        # path_save, rio_extractor, writer = args[0], args[1], args[2]
        for ifile in files:
            if ifile == 'Thumbs.db':
                continue
            img = cv2.imread(os.path.join(self._img_path, ifile), cv2.IMREAD_COLOR)
            img_crop = img[52:428, 140:, :]
            img_resize = cv2.resize(img_crop, (640, 480))
            mapper(img, img_resize, ifile, *args, **kwargs)

    def iterate(self, mapper, *args, **kwargs):
        with codecs.open(self.csv_path, 'r', 'utf-8') as csv_file:
            reader = csv.DictReader(csv_file)
            for row in reader:
                # filename = row[0]
                # x0, x2, y0, y2 = map(lambda x: float(x), row[1:])
                filename = row['file_name']
                x0 = float(row['x0'])
                x2 = float(row['x2'])
                y0 = float(row['y0'])
                y2 = float(row['y2'])
                img = cv2.imread(os.path.join(self._img_path, filename), cv2.IMREAD_GRAYSCALE)
                mapper(img, (x0, x2, y0, y2), filename, *args, **kwargs)

    def get(self, filename):
        # dir = iitd_util.get_pid(filename)
        img = cv2.imread(os.path.join(self._img_path, filename), cv2.IMREAD_COLOR)
        img_crop = img[65:535, 175:, :]
        img_resize = cv2.resize(img_crop, (640, 480))
        return img_resize


class TongjiLoader(DataLoader):
    def map(self, mapper, *args, **kwargs):
        files = os.listdir(self._img_path)
        # path_save, rio_extractor, writer = args[0], args[1], args[2]
        for ifile in files:
            if ifile == 'Thumbs.db':
                continue
            img = cv2.imread(os.path.join(self._img_path, ifile), cv2.IMREAD_COLOR)
            # img = cv2.resize(img, (640, 480))
            img = img[int(300 - 240):int(300 + 240), int(400 - 320):int(400 + 320), :]
            mapper(img, ifile, *args, **kwargs)

    def get(self, filename):
        # dir = iitd_util.get_pid(filename)
        img = cv2.imread(os.path.join(self._img_path, filename), cv2.IMREAD_COLOR)
        # img = cv2.resize(img, (640, 480))
        img = img[int(300 - 240):int(300 + 240), int(400 - 320):int(400 + 320), :]
        return img
